"""
模型预加载示例
功能：演示如何在镜像构建时预加载大型模型（如 Stable Diffusion）
特点：使用 run_function 在构建时下载模型，运行时直接使用
注意：需要配置 huggingface-secret 密钥
运行：modal run 04_model_preloading.py
"""
import modal
import os

app = modal.App("model-preloading-app")

# Hugging Face 密钥，用于下载私有模型
hf_secret = "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"

def download_model():
    # 这个函数在镜像构建时运行一次
    from diffusers import StableDiffusionPipeline
    
    model_id = "runwayml/stable-diffusion-v1-5"
    pipe = StableDiffusionPipeline.from_pretrained(
        model_id, use_auth_token=os.environ["HF_TOKEN"]
    )
    # 将模型保存到镜像内的 /model 目录
    pipe.save_pretrained("/model")

# 构建镜像时运行 download_model 函数
image = (
    modal.Image.debian_slim()
    .pip_install("diffusers[torch]", "transformers")
    .run_function(download_model, secrets=[hf_secret])
)

@app.function(image=image, gpu="A10G")
def generate_image():
    # 函数运行时，模型已经存在于 /model 目录
    from diffusers import StableDiffusionPipeline
    import torch
    
    pipe = StableDiffusionPipeline.from_pretrained(
        "/model", torch_dtype=torch.float16
    ).to("cuda")
    
    prompt = "a photo of an astronaut riding a horse on mars"
    image = pipe(prompt).images[0]
    # 返回图片数据
    return "Image generated successfully"

@app.local_entrypoint()
def main():
    result = generate_image.remote()
    print(result)